
import argparse
import os
import sys
from pathlib import Path

import torch
import torch.backends.cudnn as cudnn
from config import *

FILE = Path(__file__).resolve()
ROOT = FILE.parents[0]  # YOLOv5 root directory
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))  # add ROOT to PATH
ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative

from yolo_utils.dataloaders import letterbox
from models.common import DetectMultiBackend
from yolo_utils.general import (LOGGER, check_file, check_img_size, check_imshow, check_requirements, colorstr, cv2,
                           increment_path, non_max_suppression, print_args, scale_coords, strip_optimizer, xyxy2xywh)
from yolo_utils.plots import Annotator, colors, save_one_box
from yolo_utils.torch_utils import select_device, time_sync
from common.utils import *
from yolov5_xx.config  import get_config


@torch.no_grad()
def run_detect(
        model,  # model.pt path(s)
        imgs,  # file/dir/URL/glob, 0 for webcam
        imgsz,  # inference size (height, width)
        conf_thres=0.25,  # confidence threshold
        iou_thres=0.45,  # NMS IOU threshold
        max_det=1000,  # maximum detections per image
        device='',  # cuda device, i.e. 0 or 0,1,2,3 or cpu
        view_img=False,  # show results
        save_txt=False,  # save results to *.txt
        save_conf=False,  # save confidences in --save-txt labels
        save_crop=False,  # save cropped prediction boxes
        nosave=False,  # do not save images/videos
        classes=None,  # filter by class: --class 0, or --class 0 2 3
        agnostic_nms=False,  # class-agnostic NMS
        augment=False,  # augmented inference
        visualize=False,  # visualize features
        name='exp',  # save results to project/name
        exist_ok=False,  # existing project/name ok, do not increment
        line_thickness=3,  # bounding box thickness (pixels)
        hide_labels=False,  # hide labels
        hide_conf=False,  # hide confidences
        half=False,  # use FP16 half-precision inference
):
    if not isinstance(imgs, list):
        imgs = [imgs]

    # Load model
    stride, names, pt = model.stride, model.names, model.pt
    #print(names)

    def trans_img(img0):
        img = letterbox(img0, imgsz, stride=stride, auto=pt)[0]
        # Convert
        img = img.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB
        img = np.ascontiguousarray(img)
        return img

    dt, seen = [0.0, 0.0, 0.0], 0
    outs = []
    for im0s in imgs:
        im = trans_img(im0s)
        t1 = time_sync()
        im = torch.from_numpy(im).to(device)
        im = im.half() if model.fp16 else im.float()  # uint8 to fp16/32
        im /= 255  # 0 - 255 to 0.0 - 1.0
        if len(im.shape) == 3:
            im = im[None]  # expand for batch dim
        t2 = time_sync()
        dt[0] += t2 - t1

        # Inference
        visualize = False
        pred = model(im, augment=augment, visualize=visualize)
        t3 = time_sync()
        dt[1] += t3 - t2
        # NMS
        pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
        dt[2] += time_sync() - t3

        # Second-stage classifier (optional)
        # pred = utils.general.apply_classifier(pred, classifier_model, im, im0s)

        out = []
        im0 = im0s
        #im0 = im0s.copy()
        #annotator = Annotator(im0, line_width=line_thickness, example=str(names))
        # Process predictions
        save_img = False
        for i, det in enumerate(pred):  # per image
            seen += 1

            #save_path = str(save_dir / p.name)  # im.jpg
            # = './'
            #txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}')  # im.txt
            #s += '%gx%g ' % im.shape[2:]  # print string
            gn = torch.tensor(im0s.shape)[[1, 0, 1, 0]]  # normalization gain whwh
            #imc = im0.copy() if save_crop else im0  # for save_crop
            if len(det):
                det = det.cpu()
                # Rescale boxes from img_size to im0 size
                det[:, :4] = scale_coords(im.shape[2:], det[:, :4], im0.shape).round()
                # Print results
                #for c in det[:, -1].unique():
                #    n = (det[:, -1] == c).sum()  # detections per class
                #    #s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string

                # Write results
                for *xyxy, conf, cls in reversed(det):
                    x,y,x1,y1 = xyxy
                    c = int(cls)  # integer class
                    out.append([[float(x), float(y), float(x1), float(y1)], names[c], float(conf)])
                    if 0:
                        label = f'{names[c]} {conf:.2f}'
                        #annotator.box_label(xyxy, label, color=colors(c, True))

                    '''if save_txt:  # Write to file
                        xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
                        line = (cls, *xywh, conf) if save_conf else (cls, *xywh)  # label format
                        with open(f'{txt_path}.txt', 'a') as f:
                            f.write(('%g ' * len(line)).rstrip() % line + '\n')'''
                    #save_dir = 'D:/'
                    #if save_crop:
                    #    save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True)

                #return [[im0, []]]
                
        #im0 = annotator.result()
        outs.append([im0, out])

    #print(dt)
    # Print results
    #return [[im0, []]]
    return outs

def loadinfo(fn, model):
    jsonfn = os.path.split(fn)[0] + '/info.json'
    if os.path.exists(jsonfn):
        import json
        data = json.load(open(jsonfn, 'r'))
        model.names = data['names']
    return

def open_model(fn, exts, device, dnn, fp16):
    for ext in exts:
        weights = fn.replace('.pt', ext)
        try:
            model = DetectMultiBackend(weights, device=device, dnn=dnn, fp16=fp16)
            print(f'load {weights}')
            loadinfo(fn, model)
            return model
        except Exception as e:
            print(f'open_model??: {weights} {e}')
            continue
    return None

# ext '.engine' '.pt', '.onnx'
def get_detect_onepic(pa, xinghao, imgsz, name, exts = ['.engine', '.pt']):
    opt = get_config(pa, xinghao, imgsz, name)
    conf_thres=0.25  # confidence threshold
    iou_thres=0.45  # NMS IOU threshold
    imgsz = [opt.imgsz]*2
    weights = opt.best
    if 0:
        device = select_device('')
    else:
        device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
        #device = 'cpu'
        device = torch.device(device)
    fp16 = True
    dnn = False  # use OpenCV DNN for ONNX inference
    #exts = ['.engine', '.onnx', '.pt']
    #exts = ['.pt']
    model = open_model(opt.best, exts, device=device, dnn=dnn, fp16=fp16)
    #print(f'load {weights}')
    stride, names, pt = model.stride, model.names, model.pt
    print(names)
    imgsz = check_img_size(imgsz, s=stride)  # check image size
    bs = 1
    model.warmup(imgsz=(1 if pt else bs, 3, *imgsz))  # warmup
    def detect_onepic(img):
        outs = run_detect(model, [img], imgsz, max_det=5, device=device)
        return outs[0]
    
    return detect_onepic

# def get_detect_imgs(pa, imgsz, name, ext = '.engine'):
#     opt = get_config(pa, imgsz, name)
def get_detect_imgs(pa, xinghao,imgsz, name, ext = '.engine'):
    opt = get_config(pa, xinghao,imgsz, name)
    conf_thres=0.15  # confidence threshold
    iou_thres=0.45  # NMS IOU threshold
    imgsz = [opt.imgsz]*2
    weights = opt.best
    weights = opt.best.replace('.pt', ext)
    #weights = opt.best.replace('.pt', '.onnx')
    print(f'load {weights}')
    device = select_device('')
    half = True
    dnn=False  # use OpenCV DNN for ONNX inference
    model = DetectMultiBackend(weights, device=device, dnn=dnn, fp16=half)
    stride, names, pt = model.stride, model.names, model.pt
    imgsz = check_img_size(imgsz, s=stride)  # check image size
    bs = 1
    model.warmup(imgsz=(1 if pt else bs, 3, *imgsz))  # warmup
    def detect_imgs(imgs):
        outs = run_detect(model, imgs, imgsz, device=device)
        return outs
    
    return detect_imgs

def test_detect_onepic(pa, imgsz, name, testpa):
    import time
    detect_onepic = get_detect_onepic(pa, imgsz, name)
    os.chdir(testpa)
    mkdir('out')
    mkdir('out0')
    os.system('del out\\*.jpg')
    os.system('del out0\\*.jpg')
    li = os.listdir(testpa)

    cnt = 0
    cnt_time = 0
    for fn in li:
        fn1,ext = os.path.splitext(fn)
        if ext !='.jpg':
            continue
        path = f'{fn}'
        img0 = cv2.imdecode(np.fromfile(path, dtype=np.uint8), cv2.IMREAD_COLOR)
        t1 = time.time()
        img, out = detect_onepic(img0)
        t2 = time.time()
        cnt_time += t2-t1
        cnt+=1
        #print(out)
        print(f'{cnt_time/cnt:.5f} {fn} {len(out)}')
        #if len(out)==0:
            #cv2.imwrite(f'./out0/{fn1}_out.jpg', img)
        cv2.imwrite(f'./out/{fn1}_out.jpg', img)
    
    return 0

if __name__ == '__main__':
    pa = 'D:/data/211105宜美哲活塞环/train/imgs'
    pa = 'D:/code/git/ywlydd/deepgui/yolov5_xx/imgs/jieba'
    pa = 'E:/data/220309筷子/train'
    pa = 'E:/data/210417jieba/jieba'
    pa = 'E:/data/detect/jieba'
    pa = 'E:/data/220401小鸟/train', 416, 'xiaoniao'
    pa = 'D:/data/220309筷子/train', 416, 'kuaizi'
    pa = 'D:/data/220329竹筷/标注caise/mini5', 320, 'kuaizi'
    pa = 'D:/data/220329竹筷/标注caise/mini5', 416, 'kuaizi'
    testpa = 'D:/data/220329竹筷/train'
    testpa = pa[0]
    #testpa = 'D:/data/220309筷子/test/tt1/out'
    test_detect_onepic(*pa, testpa)
    pass


